from utils.basic_function import *
from collections import deque
from shapely.wkt import loads
from shapely.ops import snap


class SplitTool(PreprocessingTool):
    def __init__(self, table, search_threshold, split_X_threshold, preserve_col_list=None):
        super(SplitTool, self).__init__(table)
        self.search_threshold = search_threshold
        self.split_X_threshold = split_X_threshold
        self.preserve_col_list = preserve_col_list  # 分割后子路段要保留的属性

    def split_intersects_table(self):
        pk_list = pd.read_sql(f"select {self.table.pk} from {self.table.name} order by {self.table.pk} asc", conn)[self.table.pk].values.tolist()
        self.set_pk_list(deque(pk_list))
        self.split_intersects_T_workflow()

        pk_list = pd.read_sql(f"select {self.table.pk} from {self.table.name} order by {self.table.pk} asc", conn)[self.table.pk].values.tolist()
        self.set_pk_list(deque(pk_list))
        self.split_intersects_X_workflow()

    def split_intersects_T_workflow(self):
        print("线段总数：", len(self.pk_list))
        count = 0
        while len(self.pk_list) != 0:
            pk_self = self.pk_list.popleft()  # 出队
            print(pk_self)
            geo_series = self.table.select_by_attr(self.table.pk, pk_self)
            geom = geo_series[self.table.geom].iloc[0]
            st_coord = geom.coords[0]
            ed_coord = geom.coords[-1]
            self.split_intersects(st_coord, and_where=f"and {self.table.pk}!={pk_self}")
            self.split_intersects(ed_coord, and_where=f"and {self.table.pk}!={pk_self}")
            count += 1
            if count == 500:
                self.table.vacuum_analyze()
                count = 0

    def split_intersects_X_workflow(self):
        geom_col = self.table.geom
        pk_col = self.table.pk
        print("线段总数：", len(self.pk_list))
        count = 0
        while len(self.pk_list) != 0:
            pk_self = self.pk_list.popleft()  # 出队
            print(pk_self)
            gdf_cross = geopd.read_postgis(f"""
            select 
                *,
                st_AsText(
                    st_intersection(
                        {geom_col}, 
                        (select {geom_col} from {self.table.name} where {pk_col}={int(pk_self)})
                    )
                ) as cross
                from {self.table.name}
                where st_crosses(
                    {geom_col}, 
                    (select {geom_col} from {self.table.name} where {pk_col}={int(pk_self)}))
            """, conn, geom_col=geom_col)
            self.try_split_X(pk_self, gdf_cross)
            count += 1
            if count == 500:
                self.table.vacuum_analyze()
                count = 0

    def split_intersects(self, pt_coord, and_where=""):
        gdf = self.table.buffer_intersect_by_pt(pt_coord, self.search_threshold, and_where=and_where)
        split_result = self.try_split(pt_coord, gdf)
        inserted_pk_list = split_result["inserted"]
        deleted_pk_list = split_result["deleted"]
        self.pk_list.extend(inserted_pk_list)
        for pk_d in deleted_pk_list:
            try:
                self.pk_list.remove(pk_d)
            except ValueError:
                pass

    def try_split(self, pt, gdf):
        """
        尝试在pt处打断gdf中的线段，若可以打断，返回新增和删除列表，若不能打断，返回空列表
        """
        # pt_coord到线段任一端点距离都超过阈值，则打断，否则不打断。
        inserted_pk_list = []
        deleted_pk_list = []
        if gdf.shape[0] == 0:
            return {
                "deleted": [],
                "inserted": []
            }
        for i in range(gdf.shape[0]):
            line_pt_list = list(gdf[self.table.geom].iloc[i].coords)
            dist_st = cal_dist(pt, line_pt_list[0])
            dist_ed = cal_dist(pt, line_pt_list[-1])
            pk = gdf[self.table.pk].iloc[i]
            # 计算点到两个端点的距离，若都超过阈值则打断
            if dist_ed > self.search_threshold and dist_st > self.search_threshold:
                proj_result = project_pt_to_polyline(pt, line_pt_list)
                pre_node = proj_result["pre_node"]
                proj_pt = proj_result["proj_pt"]
                split_line_list = split_line_at_pt(line_pt_list, pre_node, proj_pt)
                # 事务开始
                with conn:
                    self.table.delete_by_attr(self.table.pk, pk)
                    temp = []
                    for pt_list in split_line_list:
                        geom = LineString(pt_list)
                        if self.preserve_col_list:
                            attr_dict = dict()
                            for col in self.preserve_col_list:
                                attr_dict[col] = gdf[col].iloc[i]
                        else:
                            attr_dict = dict()
                        attr_dict[self.table.geom] = geom
                        new_pk = self.table.insert(attr_dict)
                        temp.append(new_pk)
                # 事务正常结束，记录插入和删去的pk
                deleted_pk_list.append(pk)
                inserted_pk_list.extend(temp)
                print(f"split:{pk}")
        return {
            "deleted": deleted_pk_list,
            "inserted": inserted_pk_list
        }

    def try_split_X(self, pk_self, gdf_cross):
        """
        检查geom_self与gdf中的线段是否cross，若是，返回新增和删除列表，若不能打断，返回空列表
        """
        # 交点到线段任一端点距离都超过阈值，则打断，否则不打断。
        if gdf_cross.shape[0] == 0:
            return
        gdf_self = self.table.select_by_attr(self.table.pk, pk_self)
        geom_self = gdf_self[self.table.geom].iloc[0]

        # 获取pk_self要保留的属性
        if self.preserve_col_list:
            attr_dict_self = {}
            for col in self.preserve_col_list:
                attr_dict_self[col] = gdf_self[col].iloc[0]
        else:
            attr_dict_self = dict()

        # 遍历gdf_cross，判断X形交叉
        for i in range(gdf_cross.shape[0]):
            # 获取cross信息
            geom_cross = gdf_cross.iloc[i][self.table.geom]
            pk_cross = gdf_cross[self.table.pk].iloc[i]
            cross_result = loads(gdf_cross.iloc[i]["cross"])  # 获取postgis计算出的交点 避免用shapely计算相交，导致前后不一致

            # 获取pk_cross要保留的属性
            if self.preserve_col_list:
                attr_dict_cross = {}
                for col in self.preserve_col_list:
                    attr_dict_cross[col] = gdf_cross[col].iloc[i]
            else:
                attr_dict_cross = dict()

            if cross_result.type == "MultiPoint":
                for pt_num in range(len(cross_result)):
                    pt_cross = cross_result[pt_num].coords[0]
                    # 只要有一条可以打断，self的形状就改变了，因此结束循环。
                    modified = self.try_split_X_at_pt(pt_cross, pk_cross, geom_cross, attr_dict_cross,
                                                      pk_self, geom_self, attr_dict_self)
                    if modified:
                        return
            else:
                pt_cross = cross_result.coords[0]
                modified = self.try_split_X_at_pt(pt_cross, pk_cross, geom_cross, attr_dict_cross,
                                                  pk_self, geom_self, attr_dict_self)
                if modified:
                    return

    def try_split_X_at_pt(self, pt_cross, pk_cross, geom_cross, attr_dict_cross,
                          pk_self, geom_self, attr_dict_self):
        pk_col = self.table.pk
        geom_col = self.table.geom
        # 交点到cross两个端点的距离
        cross_pt_list = list(geom_cross.coords)
        dist_ends = cal_dist_end(pt_cross, cross_pt_list)

        # 交点到self两个端点的距离
        self_pt_list = list(geom_self.coords)
        dist_ends.extend(cal_dist_end(pt_cross, self_pt_list))  # 顺序：cross_st, cross_ed, self_st, self_ed

        # 计算点到两个端点的距离，若都超过阈值则打断。
        if min(dist_ends) > self.split_X_threshold:  # 两个都打断
            # 打断cross
            split_list_cross = split_line_by_project(cross_pt_list, pt_cross)
            # 打断自身
            split_list_self = split_line_by_project(self_pt_list, pt_cross)

            # 事务开始
            with conn:
                self.table.delete_by_attr(pk_col, pk_cross)
                self.table.delete_by_attr(pk_col, pk_self)
                new_pk_list = []
                for pt_list in split_list_cross:
                    geom_new = LineString(pt_list)
                    attr_dict_cross[geom_col] = geom_new
                    new_pk = self.table.insert(attr_dict_cross)
                    new_pk_list.append(new_pk)
                for pt_list in split_list_self:
                    geom_new = LineString(pt_list)
                    attr_dict_self[geom_col] = geom_new
                    new_pk = self.table.insert(attr_dict_self)
                    new_pk_list.append(new_pk)
            # 事务正常结束，记录插入和删去的pk
            self.pk_list.extend(new_pk_list)
            print(f"split:{pk_cross}, {pk_self}")
            return True

        elif len([x for x in dist_ends if self.search_threshold < x < self.split_X_threshold]) == 1:  # 打断一个，捕捉到另一个的端点上
            nearest_pt = dist_ends.index(min(dist_ends))
            with conn:
                if nearest_pt in [2, 3]:  # 打断cross，并捕捉到self上
                    split_list = split_line_by_project(cross_pt_list, pt_cross)
                    attr_dict = attr_dict_cross
                    snap_coord = self_pt_list[0 if nearest_pt == 2 else -1]
                    pk_delete = pk_cross
                    modified = False
                else:  # 打断self，捕捉到cross上
                    split_list = split_line_by_project(self_pt_list, pt_cross)
                    attr_dict = attr_dict_self
                    snap_coord = cross_pt_list[0 if nearest_pt == 0 else -1]
                    pk_delete = pk_self
                    modified = True

                self.table.delete_by_attr(pk_col, pk_delete)
                new_pk_list = []
                for pt_list in split_list:
                    geom_new = LineString(pt_list)
                    geom_new = snap(geom_new, Point(snap_coord), self.split_X_threshold)
                    attr_dict[geom_col] = geom_new
                    new_pk = self.table.insert(attr_dict)
                    new_pk_list.append(new_pk)
            # 事务正常结束，记录插入和删去的pk
            self.pk_list.extend(new_pk_list)
            print(f"split:{pk_cross}")
            return modified
        else:
            return False

    def split_by_layer(self, target_layer):
        pk_list = target_layer.buffer_intersect_by_layer(self.table)[target_layer.pk]
        for pk in pk_list:
            gdf = geopd.read_postgis(f"""select {target_layer.pk}, {target_layer.geom} from {target_layer.name} 
                where {target_layer.pk} = {pk}""", conn,
                geom_col=target_layer.geom)
            print("pk: ", pk)
            line_coord_list = list(gdf.loc[0, target_layer.geom].coords)
            st = line_coord_list[0]
            ed = line_coord_list[-1]
            self.split_intersects(st)
            self.split_intersects(ed)
